#include <iostream>
#include <vector>
#include "../../lib/binary_tree.h"

BinaryTreeNode* constructBinaryTreeCore(const std::vector<int> &pre_order,
    const std::vector<int> &in_order,
    const int &pre_order_begin,
    const int &pre_order_end,
    const int &in_order_begin,
    const int &in_order_end) {
  BinaryTreeNode* root = new BinaryTreeNode();
  root->value = pre_order[pre_order_begin];
  root->left = nullptr;
  root->right = nullptr;

  if (pre_order_begin == pre_order_end) {
    if (in_order_begin == in_order_end &&
        pre_order[pre_order_begin] == in_order[in_order_end]) {
      return root;
    } else {
      std::cout << "Invalid input." << std::endl;
    }
  }

  int root_in_order = in_order_begin;
  while (in_order[root_in_order] != pre_order[pre_order_begin] && root_in_order < in_order_end) {
    ++root_in_order;
  }
  if (in_order[root_in_order] != pre_order[pre_order_begin]) {
    std::cout << "Invalid input." << std::endl;
  }
  
  int left_num = root_in_order - in_order_begin;
  int right_num = in_order_end - root_in_order;
  if (left_num > 0) {
    root->left = constructBinaryTreeCore(pre_order, in_order, pre_order_begin + 1, pre_order_begin + left_num, in_order_begin, root_in_order - 1);
  }
  if (right_num > 0) {
    root->right = constructBinaryTreeCore(pre_order, in_order, pre_order_begin + 1 + left_num, pre_order_begin + left_num + right_num, root_in_order + 1, in_order_end);
  }
}

BinaryTreeNode* constructBinaryTree(const std::vector<int> &pre_order,
                                    const std::vector<int> &in_order) {
  int pre_order_len = pre_order.size();
  int in_order_len = in_order.size();
  if (pre_order_len != in_order_len) {
    std::cout << "Invalid input." << std::endl;
    return nullptr;
  }
  return constructBinaryTreeCore(pre_order, in_order, 0, pre_order_len - 1, 0, in_order_len - 1);
}

int main() {
  std::vector<int> pre_order = {1, 2, 4, 7, 3, 5, 6, 8};
  std::vector<int> in_order = {4, 7, 2, 1, 5, 3, 8, 6};
  auto ret = constructBinaryTree(pre_order, in_order);
  PrintBinaryTree(ret);
  DeleteBinaryTree(ret);
  return 0;
}
